#include "set.h"

#include "mem.h"
#include "assert.h"

#include <limits.h>

#define T Set_T

struct T {
    int size;
    int (*cmp)(const void *x, const void *y);
    unsigned int (*hash)(const void *x);
    int length;
    unsigned int timestamp;

    struct member {
        struct member *link;
        const void *member;
    } **buckets;
};

static int
cmpatom(const void *x, const void *y)
{
    return x != y;
}

static unsigned int
hashatom(const void *key)
{
    return (size_t)key >> 2; // address aligning, the right 2 bit may be 0
}

static void
add_to_set(T set, const void *member)
{
    struct member *p;
    int slot = set->hash(member) % set->size;

    NEW(p);
    p->member = member;
    p->link = set->buckets[slot];

    set->buckets[slot] = p;
    set->length++;
}

T Set_new(int hint, int cmp(const void *x, const void *y), unsigned int hash(const void *x))
{
    assert(hint >= 0);

    int primes[] = {509, 509, 1021, 2053, 4093, 8191, 16381, 32771, 65521, INT_MAX};
    int n;
    for (n = 1; primes[n] <= hint && primes[n] < INT_MAX; n++)
        ;
    int size = primes[n - 1];

    T set = (T)ALLOC(sizeof(*set) + size * sizeof(set->buckets[0]));
    set->size = size;
    set->cmp = cmp != NULL ? cmp : cmpatom;
    set->hash = hash != NULL ? hash : hashatom;
    set->length = 0;
    set->timestamp = 0;
    set->buckets = (struct member**)(set + 1);

    for (int i = 0; i < size; i++)
        set->buckets[i] = NULL;

    return set;
}

void Set_free(T *set)
{
    struct member *p, *q;
    if ((*set)->length > 0) {
        for (int i = 0; i < (*set)->size; i++) {
            for (p = (*set)->buckets[i]; p; p = q) {
                q = p->link;
                FREE(p);
            }
        }
    }
    FREE(*set);
}

int Set_length(T set)
{
    assert(set);
    return set->length;
}

static struct member*
search_member(T set, const void *member)
{
    int slot = set->hash(member) % set->size;
    struct member *p;
    for (p = set->buckets[slot]; p; p = p->link) {
        if (set->cmp(p->member, member) == 0)
            return p;
    }
    return NULL;
}

int Set_member(T set, const void *member)
{
    assert(set);
    assert(member);

    struct member *p = search_member(set, member);
    return p != NULL;
}

void Set_put(T set, const void *member)
{
    assert(set);
    assert(member);

    struct member *p = search_member(set, member);
    if (p == NULL) {
        add_to_set(set, member);

    } else {
        p->member = member;
    }

    set->timestamp++;
}

void* Set_remove(T set, const void *member)
{
    assert(set);
    assert(member);
    set->timestamp++;

    struct member **pp;
    int slot = set->hash(member) % set->size;
    for (pp = &set->buckets[slot]; *pp; pp = &(*pp)->link) {
        if (set->cmp((*pp)->member, member) == 0) {
            struct member *p = *pp;
            void *member = (void*)p->member;
            *pp = p->link;
            FREE(p);
            set->length--;
            return member;
        }
    }
    return NULL;
}

void Set_map(T set, void apply(const void *member, void *cl), void *cl)
{
    assert(set);
    assert(apply);

    struct member *p;
    unsigned int stamp = set->timestamp;

    for (int i = 0; i < set->size; i++) {
        for (p = set->buckets[i]; p; p = p->link) {
            apply(p->member, cl);
            assert(set->timestamp == stamp);
        }
    }
}

void** Set_toArray(T set, void *end)
{
    assert(set);

    void **array = ALLOC((set->length + 1) * sizeof(*array));
    int j = 0;
    struct member *p;

    for (int i = 0; i < set->size; i++) {
        for (p = set->buckets[i]; p; p = p->link) {
            array[j++] = (void*)p->member;
        }
    }
    array[j] = end;
    return array;
}

static T
copy_set(T t, int hint)
{
    T set = Set_new(hint, t->cmp, t->hash);
    struct member *p;
    for (int i = 0; i < t->size; i++) {
        for (p = t->buckets[i]; p; p = p->link) {
            add_to_set(set, p->member);
        }
    }
    return set;
}

// s + t
T Set_union(T s, T t)
{
    assert(s || t);

    if (s == NULL) {
        return copy_set(t, t->size);

    } else if (t == NULL) {
        return copy_set(s, s->size);

    } else {
        assert(s->cmp == t->cmp && s->hash == t->hash);

        int size = s->size > t->size ? s->size : t->size;
        T set = copy_set(s, size);

        struct member *p;
        for (int i = 0; i < t->size; i++) {
            for (p = t->buckets[i]; p; p = p->link) {
                if (!Set_member(set, p->member))
                    add_to_set(set, p->member);
            }
        }

        return set;
    }
}

// s * t
T Set_inter(T s, T t)
{
    assert(s || t);

    if (s == NULL) {
        return Set_new(t->size, t->cmp, t->hash);

    } else if (t == NULL) {
        return Set_new(s->size, s->cmp, s->hash);
    
    } else if (s->length < t->length) {
        return Set_inter(t, s);

    } else {
        assert(s->cmp == t->cmp && s->hash == t->hash);

        // t->length is less
        T set = Set_new(t->size, t->cmp, t->hash);
        struct member *p;
        for (int i = 0; i < t->size; i++) {
            for (p = t->buckets[i]; p; p = p->link) {
                if (Set_member(s, p->member))
                    add_to_set(set, p->member);
            }
        }
        return set;
    }
}

// s - t
T Set_minus(T s, T t)
{
    assert(s || t);

    if (s == NULL) {
        return Set_new(t->size, t->cmp, t->hash);
    
    } else if (t == NULL) {
        return copy_set(s, s->size);
    
    } else {
        assert(s->cmp == t->cmp && s->hash == t->hash);

        T set = Set_new(s->size, s->cmp, s->hash);
        struct member *p;
        for (int i = 0; i < s->size; i++) {
            for (p = s->buckets[i]; p; p = p->link) {
                if (!Set_member(t, p->member))
                    add_to_set(set, p->member);
            }
        }
        return set;
    }
}

// s / t means only in s or only in t
T Set_diff(T s, T t)
{
    assert(s || t);

    if (s == NULL) {
        return copy_set(t, t->size);
    
    } else if (t == NULL) {
        return copy_set(s, s->size);
    
    } else {
        assert(s->cmp == t->cmp && s->hash == t->hash);

        int size = s->size > t->size ? t->size : s->size;
        T set = Set_new(size, s->cmp, s->hash);

        struct member *p;
        for (int i = 0; i < s->size; i++) {
            for (p = s->buckets[i]; p; p = p->link) {
                if (!Set_member(t, p->member))
                    add_to_set(set, p->member);
            }
        }
        for (int i = 0; i < t->size; i++) {
            for (p = t->buckets[i]; p; p = p->link) {
                if (!Set_member(s, p->member))
                    add_to_set(set, p->member);
            }
        }
        return set;
    }
}

